import numpy as np
import matplotlib.pyplot as plt
from scipy.integrate import solve_ivp
import pickle


def xian2004(t,states,debug={}):
    x=states[0]
    xdot=states[1]

    xd=np.cos(t)
    f=0.
    # f=np.sin(t)
    K1=1
    k2=5

    e=xd-x
    ef=states[2]
    p=states[3]
    debug['xd']=xd

    
    rf=p-(k2+1)*e
    u=K1*np.sign(e+ef)-(k2+1)*rf+e
    efdot=-ef+rf
    pdot=-rf-(k2+1.)*(e+rf)+e-ef
    debug['u']=u
    debug['u2']=-(k2+1)*rf+e

    return np.array([xdot,u+f,efdot,pdot])


if __name__=="__main__":
    x0=np.array([0.,0.,0.,0.])
    xian2004(0,x0)

    sol=solve_ivp(xian2004,[0,10],x0)
    signals=[]
    for it,t in enumerate(sol.t):
        s={}
        xian2004(t,sol.y[:,it],debug=s)
        signals.append(s)

    plt.figure()
    plt.subplot(211)
    plt.plot(sol.t,sol.y[0],label='$x$')
    plt.plot(sol.t,sol.y[1],label='$\\dot{x}$')
    plt.plot(sol.t,sol.y[2],'--',label='$e_f$')
    plt.plot(sol.t,sol.y[3],'--',label='$p$')
    plt.plot(sol.t,[s['xd'] for s in signals],'--',label='xd')
    plt.legend()
    plt.subplot(212)
    plt.plot(sol.t,[s['u'] for s in signals],'-',label='u')
    # plt.plot(sol.t,[s['u2'] for s in signals],'-',label='u2')
    plt.legend()
    plt.show()
